"""Module for simulating radio telescope observations."""

from astropy.coordinates import SkyCoord, EarthLocation, AltAz, ICRS
from astropy.time import Time
import astropy.units as u
import numpy as np


class ObservationSimulator:
    """Base class for simulating observations."""
    
    def __init__(self):
        pass
        
    def simulate_drift_scan(self, sky_image, telescope_location, start_time, duration, azimuth=0*u.deg, elevation=90*u.deg, interval=1*u.second):
        """
        Simulate drift scan observation of a sky image.
        
        Args:
            sky_image: SkyImage object to observe
            telescope_location: EarthLocation of telescope
            start_time: Time to start observation
            duration: Duration of observation in seconds
            azimuth: Azimuth of telescope (degrees)
            elevation: Elevation of telescope (degrees)
            interval: Time interval between observations (astropy Quantity)
            
        Returns:
            numpy.ndarray: Array of intensity values over time
        """
        import psutil
        import warnings
        
        # Generate time steps with specified interval
        num_steps = int(np.ceil(duration / interval.to(u.second).value))
        times = start_time + np.arange(0, num_steps) * interval
        
        # Test run to determine the shape of returned intensity data
        telescope_direction = AltAz(alt=elevation, az=azimuth, obstime=times[0], location=telescope_location)
        celestial_coord = telescope_direction.transform_to(ICRS())
        test_intensity = sky_image.get_intensity_at_coordinate(celestial_coord)
        
        # Estimate memory usage (rough estimate)
        # Assuming each intensity value is a float64 (8 bytes)
        estimated_element_size = 8  # bytes per float64
        estimated_elements = num_steps
        # If we have frequency data or other multi-dimensional data
        if hasattr(test_intensity, 'shape') and test_intensity.shape:
            intensity_shape = test_intensity.shape
            estimated_elements *= np.prod(intensity_shape)
        
        estimated_memory_bytes = estimated_elements * estimated_element_size
        estimated_memory_mb = estimated_memory_bytes / (1024 * 1024)
        
        # Check available memory
        available_memory_mb = psutil.virtual_memory().available / (1024 * 1024)
        
        # Warn if memory usage is high
        if estimated_memory_mb > 0.5 * available_memory_mb:
            warning_msg = f"Estimated memory usage ({estimated_memory_mb:.2f} MB) exceeds 50% of available memory ({available_memory_mb:.2f} MB)."
            warnings.warn(warning_msg, ResourceWarning)
            
        # If memory usage is too high, raise an error
        if estimated_memory_mb > 0.8 * available_memory_mb:
            error_msg = f"Estimated memory usage ({estimated_memory_mb:.2f} MB) exceeds 80% of available memory ({available_memory_mb:.2f} MB). Aborting simulation."
            raise MemoryError(error_msg)
        
        # Pre-allocate array with proper shape
        if hasattr(test_intensity, 'shape') and test_intensity.shape:
            intensities = np.empty((num_steps,) + test_intensity.shape)
        else:
            intensities = np.empty(num_steps)
        
        # Set first intensity value from test run
        intensities[0] = test_intensity
        
        # For each time step, calculate telescope coordinates and extract intensity
        for i, t in enumerate(times[1:], 1):
            # Calculate telescope coordinates at this time
            telescope_direction = AltAz(alt=elevation, az=azimuth, obstime=t, location=telescope_location)
            
            # Convert to celestial coordinates (RA, Dec)
            celestial_coord = telescope_direction.transform_to(ICRS())
            
            # Extract intensity at this coordinate from sky_image
            # This will call the sky_image module's interpolation functions as needed
            intensity = sky_image.get_intensity_at_coordinate(celestial_coord)
            intensities[i] = intensity
            
        return intensities
        
    def simulate_observation(self, sky_image, telescope_params):
        """Simulate observation of a sky image with given telescope parameters."""
        pass
        
    def generate_spectrum(self, observation_data):
        """Generate spectral data from observation."""
        pass
        
    def add_noise(self, data, noise_params):
        """Add noise to simulated data."""
        pass